package ai.george.quickstart.controller;


import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatOptions;
import io.swagger.v3.oas.annotations.Operation;
import io.swagger.v3.oas.annotations.Parameter;
import io.swagger.v3.oas.annotations.tags.Tag;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.ai.chat.prompt.SystemPromptTemplate;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.util.List;
import java.util.Map;


@Tag(name = "ChatModel示例")
@RestController
@RequestMapping("/chat-model")
public class ChatModelController {


    @Autowired
    private ChatModel chatModel;


    @Operation(summary = "ChatModel-1-String文本", description = "使用的是chatModel.call(String message)调用方式")
    @GetMapping("/text")
    public String text(@RequestParam(value = "query", defaultValue = "你好，很高兴认识你，能简单介绍一下自己吗？") @Parameter(name = "query", description = "请求内容") String query) {
        return chatModel.call(query);
    }


    @Operation(summary = "ChatModel-2-Message消息", description = "使用的是chatModel.call(Message message)调用方式")
    @GetMapping("/messages")
    public String message(@RequestParam(value = "query", defaultValue = "你好，很高兴认识你，能简单介绍一下自己吗？") @Parameter(name = "query", description = "请求内容") String query) {
        SystemMessage systemMessage = new SystemMessage("你一个博学的智能聊天助手，请根据用户提问回答！");
        UserMessage userMessage = new UserMessage(query);
        return chatModel.call(systemMessage, userMessage);
    }


    @Operation(summary = "ChatModel-3-Prompt提示词", description = "自定义封装Prompt对象，使用的是chatModel.call(Prompt prompt)调用方式")
    @GetMapping("/prompt")
    public ChatResponse prompt(@RequestParam(value = "model", defaultValue = "qwen-flash") @Parameter(name = "model", description = "大模型") String model,
                               @RequestParam(value = "query", defaultValue = "你好，很高兴认识你，能简单介绍一下自己吗？") @Parameter(name = "query", description = "请求内容") String query) {
        SystemMessage systemMessage = new SystemMessage("你一个博學的智能聊天助手，请根据用户提问回答！");
        UserMessage userMessage = new UserMessage(query);
        ChatOptions options = DashScopeChatOptions.builder()
                .withModel(model)
                .withTemperature(0.7)
                .withMaxToken(1024)
                .build();
        Prompt prompt = new Prompt(List.of(systemMessage, userMessage), options);
        ChatResponse response = chatModel.call(prompt);
        return response;
    }


    @Operation(summary = "ChatModel-4-Stream流式-String文本", description = "流式调用，使用的是chatModel.stream(String msg)调用方式")
    @GetMapping("/text-stream")
    public Flux<String> textStream(@RequestParam(value = "query", defaultValue = "你好，很高兴认识你，能简单介绍一下自己吗？") @Parameter(name = "query", description = "请求内容") String query) {
        return chatModel.stream(query);
    }


    @Operation(summary = "ChatModel-5-Stream流式-Prompt提示词", description = "流式调用，使用的是chatModel.stream(Prompt prompt)调用方式")
    @GetMapping("/prompt-stream")
    public Flux<ChatResponse> promptStream(@RequestParam(value = "model", defaultValue = "qwen-flash") @Parameter(name = "model", description = "大模型") String model,
                                           @RequestParam(value = "query", defaultValue = "你好，很高兴认识你，能简单介绍一下自己吗？") @Parameter(name = "query", description = "请求内容") String query) {
        SystemMessage systemMessage = new SystemMessage("你一个博學的智能聊天助手，请根据用户提问回答！");
        UserMessage userMessage = new UserMessage(query);
        ChatOptions options = DashScopeChatOptions.builder()
                .withModel(model)
                .withTemperature(0.7)
                .withMaxToken(1024)
                .build();
        Prompt prompt = new Prompt(List.of(systemMessage, userMessage), options);
        Flux<ChatResponse> response = chatModel.stream(prompt);
        return response;
    }


    /**
     * 测试PromptTemplate
     *
     * @param args
     */
    public static void main(String[] args) {
        PromptTemplate promptTemplate = new PromptTemplate("帮我用{language}写一个{sort}排序");
        Message message = promptTemplate.createMessage(Map.of("language", "java", "sort", "冒泡"));
        System.out.println(message);

        SystemPromptTemplate systemPromptTemplate = new SystemPromptTemplate("你一个{role}的小助手，请根据用户提问回答！");
        Message systemMessage = systemPromptTemplate.createMessage(Map.of("role", "博学"));
        System.out.println(systemMessage);
    }
}
